#
#   Richardson_RK4_SHO.py
#
#   Solves the Phase space dynamics of Simple Harmonic Oscillator
#

#
#   Copyright (C) 2012
#
#   This program is free software: you can redistribute it and/or modify
#   it under the terms of the GNU General Public License as published by
#   the Free Software Foundation, either version 3 of the License, or
#   (at your option) any later version.
#
#   This program is distributed in the hope that it will be useful,
#   but WITHOUT ANY WARRANTY; without even the implied warranty of
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#   GNU General Public License for more details.
#
#   You should have received a copy of the GNU General Public License
#   along with this program.  If not, see <http://www.gnu.org/licenses/>.
#

from FunctionalBohm import SchrodingerNDT

# Problem setup

from HarmonicOscillatorND import Psi as wvfn
from HarmonicOscillatorND import V
from HarmonicOscillatorND import mass
from HarmonicOscillatorND import hbar
from HarmonicOscillatorND import all_omega

# get the minimum period with the least common multiple

def lcm(args):
    def helper(x,y):
        tmp=x
        while (tmp%y) > 1e-5: # should be == 0, but binary representation can be inexact
            tmp+=x
        return tmp
    return reduce(helper,args)

print(lcm(all_omega))

# Initial Conditions
from math import pi
timesteps = 50
T_initial = 0.0 # 0.5*pi/omega
T_final = T_initial + 2*(2*pi/lcm(all_omega)) # one period

def segments(min,max,steps):
    return [i/float(steps)*(max-min)+min for i in range(steps+1)]

def csv(*args):
    if len(args) == 1:
        return str(args[0])+"\n"
    else:
        return str(args[0])+","+csv(*args[1:])

# Use RK4 with Richardson interpolation to propagate statistical ensembles:

S = SchrodingerNDT(V,wvfn,mass,dimensions=2,dx=5e-10)

for X_ic in segments(-15,15,5): # different starting places
    for Y_ic in segments(-15,15,5): # different starting places
        print(X_ic,Y_ic)
        stepped = S.get_Trajectory_in_N_dimensions([X_ic,Y_ic],T_initial,T_final,timesteps) # trajectory as a list.
        print(stepped)